#!/bin/bash

set -e

# Source common functions and variables
source ./common.sh

MOUNT_POINTS=("proc" "sys" "dev" "dev/pts")
UMOUNT_POINTS=("dev/pts" "dev" "sys" "proc")

mount_filesystem() {
    local src="/$1"
    local dest="$ROOTFS_DIR/$1"

    case "$1" in
        proc)
            sudo mount -t proc $src "$dest"
            log_info "sudo mount -t proc $src "$dest""
            ;;
        sys)
            sudo mount -t sysfs $src "$dest"
            log_info "sudo mount -t sysfs $src "$dest""
            ;;
        dev)
            #sudo mount -t tmpfs -o mode=0755 none "$dest"
            #sudo mkdir -p "$dest/pts" "$dest/shm"
            #sudo mknod -m 666 "$dest/null" c 1 3
            #sudo mknod -m 666 "$dest/zero" c 1 5
            #sudo mknod -m 666 "$dest/random" c 1 8
            #sudo mknod -m 666 "$dest/urandom" c 1 9
            sudo mount -o bind $src "$dest"
            log_info "sudo mount -o bind $src "$dest""
            ;;
        dev/pts)
            #sudo mount -t devpts -o gid=5,mode=620 none "$dest"
            sudo mount -o bind $src "$dest"
            log_info "sudo mount -o bind $src "$dest""
            ;;
    esac
}

unmount_filesystem() {
    local dest="$ROOTFS_DIR/$1"
    if mountpoint -q "$dest"; then
        sudo umount -l "$dest"
        log_info "Unmounted $dest"
    fi
}

mount_all() {
    log_info "Mounting all filesystems"
    for fs in "${MOUNT_POINTS[@]}"; do
        if ! mountpoint -q "$ROOTFS_DIR/$fs"; then
            mount_filesystem "$fs"
        else
            log_info "$ROOTFS_DIR/$fs is already mounted"
        fi
    done
}

unmount_all() {
    log_info "Unmounting all filesystems"
    for fs in "${UMOUNT_POINTS[@]}"; do
        unmount_filesystem "$fs"
    done
}

clean_all() {
    log_info "Clean all filesystems"
    unmount_all
    sudo rm -fr "$ROOTFS_DIR"

    # Delete the old image
    rm -f "$WORK_DIR/ubuntu_rootfs.img"

    # Make sure the mount dir OK
    if mountpoint -q "$WORK_DIR/$FORGE_DIR"; then
        sudo umount -l "$WORK_DIR/$FORGE_DIR"
        log_info "Unmounted $WORK_DIR/$FORGE_DIR"
    fi

    # Remove the old directories
    sudo rm -fr "$WORK_DIR/$FORGE_DIR"
}

mount_show() {
    log_info "showing mounted filesystems"
    mount | grep $ROOTFS_DIR
}

chroot_exec() {
    log_info "Executing in chroot"
    #sudo chroot "$ROOTFS_DIR" /bin/bash -c "$1"
    sudo nsenter --target 1 --mount --uts --ipc --net --pid -- chroot "$ROOTFS_DIR" /bin/bash -c "$1"
}

enter_chroot() {
    log_info "Enter chroot environment"
    mount_all
    sudo nsenter --target 1 --mount --uts --ipc --net --pid -- chroot $ROOTFS_DIR
}

customize_system() {
    log_info "Starting system customization"
    mount_all
    chroot_exec "
    set -e
    export LC_ALL=C
    export DEBIAN_FRONTEND=noninteractive

    # Ensure /dev/null has correct permissions
    chmod 666 /dev/null

    # Ensure /tmp has correct permissions
    chmod 1777 /tmp/

    # Update package lists
    apt-get update || { echo 'Failed to update package lists'; exit 1; }

    # Install packages
    apt-get install -y \
        sudo vim xterm udhcpc ethtool ntp usbutils alsa-utils net-tools netplan.io sysstat \
        language-pack-en-base language-pack-zh-han* iputils-ping openssh-server dhcpcd5 python3-pip \
        openssh-client libmtp9 tzdata locales udev || { echo 'Failed to install packages'; exit 1; }
    echo \"-----packages installed-----\"

    echo \"-----install python3 pip tools -----\"
    python3 -m pip install --upgrade pip
    python3 -m pip install openpyxl
    python3 -m pip install --user openpyxl

    # Set timezone
    if [ ! -L /etc/localtime ] || [ ! \"\$(readlink /etc/localtime)\" = \"/usr/share/zoneinfo/$TIMEZONE\" ]; then
        ln -fs /usr/share/zoneinfo/$TIMEZONE /etc/localtime
        dpkg-reconfigure -f noninteractive tzdata
        echo \"-----timezone set-----\"
    else
        echo \"-----timezone already set-----\"
    fi

    # Set locale
    if ! locale -a | grep -q \"$LOCALE\"; then
        locale-gen $LOCALE
        update-locale LANG=$LOCALE
        echo \"-----locale set-----\"
    else
        echo \"-----locale already set-----\"
    fi

    # Add user if not exists
    if ! id -u $USERNAME >/dev/null 2>&1; then
        adduser --gecos '' --disabled-password $USERNAME
        adduser $USERNAME sudo
        echo '$USERNAME:$PASSWORD' | chpasswd
        echo \"-----user $USERNAME added-----\"
    else
        echo \"-----user $USERNAME already exists-----\"
        # Update password if needed
        echo '$USERNAME:$PASSWORD' | chpasswd
    fi

    # Desktop environment setup (if required)
    if [ \"$INSTALL_DESKTOP\" = true ]; then
        if ! dpkg -s $DESKTOP_ENV >/dev/null 2>&1; then
            apt-get install -y $DESKTOP_ENV || { echo 'Failed to install desktop environment'; exit 1; }
        fi
        if [ -f /etc/lightdm/lightdm.conf ]; then
            sed -i 's/^#autologin-user=.*/autologin-user=$USERNAME/' /etc/lightdm/lightdm.conf
            sed -i 's/^#autologin-user-timeout=.*/autologin-user-timeout=0/' /etc/lightdm/lightdm.conf
        fi
        echo \"-----desktop environment configured-----\"
    fi

    # Root login configuration (if required)
    if [ \"$ROOT_LOGIN\" = true ] && [ -f /etc/gdm3/custom.conf ]; then
        sed -i 's/^#\s*AllowRoot\s*=\s*false/AllowRoot = true/' /etc/gdm3/custom.conf
        echo \"-----root login enabled-----\"
    fi

    # Partition resize script (if required)
    if [ -n \"$RESIZE_PARTITION\" ] && [ ! -f /etc/init.d/resize2fs.sh ]; then
        cat << 'EOF' > /etc/init.d/resize2fs.sh
#!/bin/bash -e
if [ ! -e \"/usr/local/boot_flag\" ] ; then
  echo \"Resizing $RESIZE_PARTITION...\"
  resize2fs $RESIZE_PARTITION
  touch /usr/local/boot_flag
fi
EOF
        chmod +x /etc/init.d/resize2fs.sh
        cat << 'EOF' > /lib/systemd/system/resize2fs.service
[Unit]
Description=Resize filesystem on first boot
Before=lightdm.service
After=resize-helper.service

[Service]
Type=oneshot
ExecStart=/etc/init.d/resize2fs.sh
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
EOF
        systemctl enable resize2fs.service
        echo \"-----partition resize script configured-----\"
    fi

    # disable default wpa_supplicant
    systemctl disable wpa_supplicant.service
    systemctl mask wpa_supplicant.service
    echo \"-----disabled original wpa_supplicant service-----\"

    cat << EOF > /etc/systemd/system/wifibt-init.service
[Unit]
Description=WiFi and Bluetooth Initialization Service
Before=network.target
Conflicts=wpa_supplicant.service

[Service]
Type=oneshot
ExecStart=/usr/bin/wifibt-init.sh start
ExecStop=/usr/bin/wifibt-init.sh stop
RemainAfterExit=yes
User=root

[Install]
WantedBy=multi-user.target
EOF
    systemctl enable wifibt-init.service

    # customize environment
    echo 'export LS_COLORS=\"di=34:ln=36:so=35:pi=33:ex=32:bd=40;33:cd=40;33:su=37;41:sg=30;43:tw=30;42:ow=34;42:st=37;44:mi=05;37;41:or=31;41:*.tar=31:*.gz=31:*.zip=31:*.tgz=31:*.jpg=35:*.png=35:\"' >> /home/ubuntu/.bashrc

    # customize vim
    cat << EOF > /home/ubuntu/.vimrc
set nu
set ru
set smartindent
set autoindent
set cindent
\"set cino=:0g0t0(sus
:set cinoptions=l1

set et
set shiftwidth=4
set tabstop=4
set softtabstop=4
set clipboard+=unnamed

set showmatch
set hlsearch
set incsearch
\"set et
set nocp
set ai
set ignorecase
\"set mouse=a

highlight WhitespaceEOL ctermbg=red guibg=red
match WhitespaceEOL /\s\+$/

set nocompatible
set backspace=indent,eol,start
set ruler
set showcmd
set wildmenu
set fo=cqrt
set laststatus=2
set textwidth=78
set ww=<,>,h,l
\" set no error bells
set noeb visualbell
let mapleader = \",\"
syn on
syntax on
filetype on
filetype plugin on
filetype indent on

set backspace=indent,eol,start

filetype plugin on
set fencs=ucs-bom,utf-8,shift-jis,gb18030,gbk,gb2312,cp936
set fencs=cp936,ucs-bom,utf-8,shift-jis,gb18030,gbk,gb2312,euc-jp
set fenc=utf-8 \"default coding

set encoding=utf-8
set termencoding=utf-8
\"set fileencodings=utf-8

autocmd VimLeave * call system('printf \"\033[0;0H\033[2J\"')

let g:AutoComplPop_IgnoreCaseOption=1
let g:AutoComplPop_MappingDriven=1

set ignorecase
EOF

    apt-get clean
    echo \"-----system customization completed-----\"
    "
}

prepare_ubuntu_base() {
    # Create work directory
    mkdir -p "$WORK_DIR"
    cd "$WORK_DIR"
    log_info "Changed to working directory: $WORK_DIR"

    # Extract the rootfs if not already extracted
    if [ ! -d "$ROOTFS_DIR/etc" ]; then
        # Download Ubuntu base if not already present
        mkdir -p $ROOTFS_DIR
        if [ ! -f "$ROOTFS_DIR/../ubuntu-base-${UBUNTU_VERSION}-base-arm64.tar.gz" ]; then
            log_info "Downloading Ubuntu base tarball"
            wget -P "$ROOTFS_DIR/.." "$UBUNTU_BASE_URL" || {
                log_error "Failed to download Ubuntu base tarball"
                            exit 1
                        }
                        log_info "Download completed"
                    else
                        log_warning "Ubuntu base tarball already exists. Skipping download."
        fi

        log_info "Extracting rootfs"
        tar xzf "$ROOTFS_DIR/../ubuntu-base-${UBUNTU_VERSION}-base-arm64.tar.gz" -C "$ROOTFS_DIR" || {
            log_error "Failed to extract rootfs"
            exit 1
        }
        log_info "Extraction completed"

        # Update sources.list
        log_info "Updating sources.list"
        cat <<EOF | sudo tee -a "$ROOTFS_DIR/etc/apt/sources.list"

deb http://mirrors.aliyun.com/ubuntu-ports/ focal main restricted
deb http://mirrors.aliyun.com/ubuntu-ports/ focal-updates main restricted
deb http://mirrors.aliyun.com/ubuntu-ports/ focal universe
deb http://mirrors.aliyun.com/ubuntu-ports/ focal-updates universe
deb http://mirrors.aliyun.com/ubuntu-ports/ focal multiverse
deb http://mirrors.aliyun.com/ubuntu-ports/ focal-updates multiverse
deb http://mirrors.aliyun.com/ubuntu-ports/ focal-backports main restricted universe multiverse
deb http://mirrors.aliyun.com/ubuntu-ports/ focal-security main restricted
deb http://mirrors.aliyun.com/ubuntu-ports/ focal-security universe
deb http://mirrors.aliyun.com/ubuntu-ports/ focal-security multiverse
EOF
        log_info "sources.list updated"
    else
        log_warning "Rootfs already extracted. Skipping extraction."
    fi

    # Install qemu-user-static if not already installed
    if ! command -v qemu-aarch64-static &> /dev/null; then
        log_info "Installing qemu-user-static"
        sudo apt-get update && sudo apt-get install -y qemu-user-static || {
            log_error "Failed to install qemu-user-static"
                    exit 1
                }
                log_info "qemu-user-static installed"
            else
                log_info "qemu-user-static already installed"
    fi

    # Copy qemu-aarch64-static to rootfs if not already present
    if [ ! -f "$ROOTFS_DIR/usr/bin/qemu-aarch64-static" ]; then
        log_info "Copying qemu-aarch64-static to rootfs"
        sudo cp "$(which qemu-aarch64-static)" "$ROOTFS_DIR/usr/bin/" || {
            log_error "Failed to copy qemu-aarch64-static to rootfs"
                    exit 1
                }
                log_info "qemu-aarch64-static copied"
            else
                log_info "qemu-aarch64-static already present in rootfs"
    fi

    # Copy resolv.conf
    SRC_MD5=$(md5sum /etc/resolv.conf | awk '{print $1}')
    DST_MD5=$(md5sum "$ROOTFS_DIR/etc/resolv.conf" 2>/dev/null | awk '{print $1}')
    log_info "src md5: $SRC_MD5, dest md5: $DST_MD5"
    if [ "$SRC_MD5" != "$DST_MD5" ]; then
        log_info "Copying /etc/resolv.conf to rootfs"
        sudo cp /etc/resolv.conf "$ROOTFS_DIR/etc/" || {
            log_error "Failed to copy /etc/resolv.conf to rootfs"
                    exit 1
                }
                log_info "/etc/resolv.conf copied"
            else
                log_info "/etc/resolv.conf already up to date in rootfs"
    fi

    # Remove the base tarball if exist
    if [ -f "$ROOTFS_DIR/../ubuntu-base-${UBUNTU_VERSION}-base-arm64.tar.gz" ]; then
        echo "remove $ROOTFS_DIR/../ubuntu-base-${UBUNTU_VERSION}-base-arm64.tar.gz"
        #rm -fr "$ROOTFS_DIR/../ubuntu-base-${UBUNTU_VERSION}-base-arm64.tar.gz"
    fi

    log_info "Base Ubuntu image prepared"
}

finalize() {
    unmount_all

    # Delete the old image
    rm -f "$WORK_DIR/ubuntu_rootfs.img"

    # Make sure the mount dir OK
    if mountpoint -q "$WORK_DIR/$FORGE_DIR"; then
        sudo umount -l "$WORK_DIR/$FORGE_DIR"
        log_info "Unmounted $WORK_DIR/$FORGE_DIR"
    fi

    # Remove the old directories
    sudo rm -fr "$WORK_DIR/$FORGE_DIR"

    # Create image file
    log_info "Creating fw: ubuntu_rootfs.img"
    dd if=/dev/zero of="$WORK_DIR/ubuntu_rootfs.img" bs=1M count=$IMAGE_SIZE || {
        log_error "Failed to create image file"
            exit 1
    }
    log_info "Image file created"

    # Format the image
    log_info "Formatting image"
    mkfs.ext4 "$WORK_DIR/ubuntu_rootfs.img" || {
        log_error "Failed to format image"
        exit 1
    }
    log_info "Image formatted"

    # Mount image and copy files
    mkdir -p "$WORK_DIR/$FORGE_DIR"
    sudo mount "$WORK_DIR/ubuntu_rootfs.img" "$WORK_DIR/$FORGE_DIR" || {
        log_error "Failed to mount image"
        exit 1
    }
    log_info "Mounting image and copying files"

    sudo cp -a "$ROOTFS_DIR"/* "$WORK_DIR/$FORGE_DIR/" || {
        log_error "Failed to copy files to image"
        umount "$WORK_DIR/$FORGE_DIR"
        exit 1
    }
    sudo cp -a "$WORK_DIR/rk3588_files"/wifi* "$WORK_DIR/$FORGE_DIR/usr/bin/" || {
        log_error "Failed to copy wifi manage scripts"
    }
    sudo cp -a "$WORK_DIR/rk3588_files"/wpa_supplicant.conf "$WORK_DIR/$FORGE_DIR/etc/" || {
        log_error "Failed to copy wpa_supplicant.conf"
    }
    sudo cp -a "$WORK_DIR/rk3588_files"/brcm_patchram_plus1 "$WORK_DIR/$FORGE_DIR/usr/bin/" || {
        log_error "Failed to copy brcm_patchram_plus1"
    }

    sudo umount -l "$WORK_DIR/$FORGE_DIR"
    log_info "Files copied and image unmounted"

    # Fix and optimize image
    log_info "Fixing and optimizing image"
    e2fsck -f -y "$WORK_DIR/ubuntu_rootfs.img" || log_warning "e2fsck reported errors"
    resize2fs -M "$WORK_DIR/ubuntu_rootfs.img" || log_warning "resize2fs failed"
    log_info "Image fixed and optimized"

    log_info "Ubuntu ARM image creation complete. Image size:"
    du -sh "$WORK_DIR/ubuntu_rootfs.img"

    cp "$WORK_DIR/ubuntu_rootfs.img" /workspace/files/firmware/ || {
        log_error "ubuntu_rootfs.img deployment finished!"
    }

    rm -fr "$WORK_DIR/$FORGE_DIR"

    log_info "Image creation & deployment completed successfully!"
}

main() {
    local OPTIND
    local option
    while getopts ":pmecualsf" option; do
        case $option in
            p)
                # prepare
                prepare_ubuntu_base
                exit 0
                ;;
            m)
                # mount
                mount_all
                exit 0
                ;;
            e)
                # enter
                enter_chroot
                exit 0
                ;;
            c)
                # customize
                customize_system
                exit 0
                ;;
            u)
                # unmount
                unmount_all
                exit 0
                ;;
            a)
                # do all
                prepare_ubuntu_base
                customize_system
                unmount_all
                ;;
            l)
                # clean all
                clean_all
                ;;
            s)
                # show mounts
                mount_show
                ;;
            f)
                # finalization
                finalize
                ;;
            \?)
                log_error "Invalid option: -$OPTARG"
                exit 1
                ;;
        esac
    done
}

main "$@"
